{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/syao/.conda/envs/wisp/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import kaolin.ops.spc as spc_ops\n",
    "\n",
    "from collections import deque\n",
    "from termcolor import colored\n",
    "\n",
    "def create_dense_octree(level):\n",
    "    \"\"\"Creates a dense SPC model\n",
    "\n",
    "    Args:\n",
    "        level (int): The level at which the octree will be initialized to.\n",
    "\n",
    "    Returns:\n",
    "        (torch.ByteTensor): the octree tensor\n",
    "    \"\"\"\n",
    "    coords = np.arange(2**level)\n",
    "    points = np.array(np.meshgrid(coords, coords, coords, indexing='xy'))\n",
    "    points = points.transpose(3,2,1,0).reshape(-1, 3)\n",
    "    points = torch.from_numpy(points).short().cuda()\n",
    "    octree = spc_ops.unbatched_points_to_octree(points, level)\n",
    "    return octree\n",
    "\n",
    "def describe_tensor(torch_tensor, tensor_label, with_shape, with_content):\n",
    "    if with_shape:\n",
    "        print(f'\"{tensor_label}\" is a {torch_tensor.dtype} tensor of size {tuple(torch_tensor.shape)}')\n",
    "    if with_content:\n",
    "        print(f'Raw Content: \\n{torch_tensor.cpu().numpy()}')\n",
    "\n",
    "def color_by_level(level):\n",
    "    _colormap = [\"red\", \"blue\", \"green\", \"yellow\", \"magenta\", \"cyan\", \"grey\"]\n",
    "    return _colormap[level % len(_colormap)]\n",
    "\n",
    "def push_pop_octree(q, oct_item):\n",
    "    prefix = q.popleft()\n",
    "    bit_idx = 0\n",
    "    parsed_bits = oct_item.item()\n",
    "    while parsed_bits:\n",
    "        bit_idx += 1\n",
    "        if parsed_bits & 1:\n",
    "            if len(prefix) > 0:\n",
    "                q.append(prefix + f'-{bit_idx}')\n",
    "            else:\n",
    "                q.append(prefix + f'{bit_idx}')\n",
    "        parsed_bits >>= 1\n",
    "    return prefix\n",
    "\n",
    "def format_octree_str(octree_byte, octree_path, level_idx, max_level):\n",
    "    text = []\n",
    "\n",
    "    level_color = color_by_level(level_idx - 1)\n",
    "    text += ['Level ' + colored(f'#{level_idx}, ', level_color)]\n",
    "\n",
    "    colored_path = []\n",
    "    for i in range(len(octree_path)):\n",
    "        level_color = color_by_level(i // 2)\n",
    "        if i % 2 == 0:\n",
    "            colored_path += [colored(octree_path[i], level_color)]\n",
    "        else:\n",
    "            colored_path += [octree_path[i]]\n",
    "    colored_path = ''.join(colored_path)\n",
    "    text += [f'Path{colored_path},    ']\n",
    "    text += [' ' for _ in range((max_level - level_idx) * 2)]\n",
    "\n",
    "    text += ['{0:08b}'.format(octree_byte)]\n",
    "\n",
    "    return ''.join(text)\n",
    "\n",
    "def describe_octree(octree, level, limit_levels=None):\n",
    "    bit_counter = lambda x: bin(x).count('1')\n",
    "    level_idx, curr_level_remaining_cells, next_level_cells = 1, 1, 0\n",
    "    octree_paths = deque('*')\n",
    "\n",
    "    for oct_idx, octree_byte in enumerate(octree):\n",
    "\n",
    "        octree_path = push_pop_octree(octree_paths, octree_byte)\n",
    "        if limit_levels is None or level_idx in limit_levels:\n",
    "            print(format_octree_str(octree_byte, octree_path, level_idx, level))\n",
    "        curr_level_remaining_cells -= 1\n",
    "        next_level_cells += bit_counter(octree_byte)\n",
    "\n",
    "        if not curr_level_remaining_cells:\n",
    "            level_idx += 1\n",
    "            curr_level_remaining_cells = next_level_cells\n",
    "            next_level_cells = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"octree\" is a torch.uint8 tensor of size (9,)\n",
      "Raw Content: \n",
      "[255 255 255 255 255 255 255 255 255]\n",
      "Level #1, Path*,      11111111\n",
      "Level #2, Path*-1,    11111111\n",
      "Level #2, Path*-2,    11111111\n",
      "Level #2, Path*-3,    11111111\n",
      "Level #2, Path*-4,    11111111\n",
      "Level #2, Path*-5,    11111111\n",
      "Level #2, Path*-6,    11111111\n",
      "Level #2, Path*-7,    11111111\n",
      "Level #2, Path*-8,    11111111\n"
     ]
    }
   ],
   "source": [
    "blas_level=2\n",
    "octree = create_dense_octree(blas_level)\n",
    "describe_tensor(torch_tensor=octree, tensor_label='octree', with_shape=True, with_content=True)\n",
    "describe_octree(octree, blas_level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"octree\" is a torch.uint8 tensor of size (4,)\n",
      "Raw Content: \n",
      "[208 128  16 128]\n",
      "Level #1, Path*,      11010000\n",
      "Level #2, Path*-5,    10000000\n",
      "Level #2, Path*-7,    00010000\n",
      "Level #2, Path*-8,    10000000\n"
     ]
    }
   ],
   "source": [
    "blas_level=2\n",
    "points = torch.tensor([[3,2,0],[3,1,1],[3,3,3]], device='cuda', dtype=torch.short)\n",
    "octree=None\n",
    "octree = spc_ops.unbatched_points_to_octree(points, blas_level)\n",
    "describe_tensor(torch_tensor=octree, tensor_label='octree', with_shape=True, with_content=True)\n",
    "describe_octree(octree, blas_level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([6], device='cuda:0')\n",
      "tensor([[0, 3, 6]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "lengths = torch.tensor([len(octree)], dtype=torch.int32)\n",
    "_, pyramid_orig, prefix = spc_ops.scan_octrees(octree, lengths)\n",
    "points = spc_ops.generate_points(octree, pyramid_orig, prefix)\n",
    "query_coords = torch.tensor([[3,2,0]], device='cuda', dtype=torch.short) #[0,2,5]\n",
    "query_coords = torch.tensor([[3,1,1]], device='cuda', dtype=torch.short) #[0,1,4]\n",
    "query_coords = torch.tensor([[3,3,3]], device='cuda', dtype=torch.short) #[0,3,6]\n",
    "\n",
    "pidx=spc_ops.unbatched_query(octree, prefix, query_coords, blas_level, with_parents=False)\n",
    "print(pidx)\n",
    "pidx=spc_ops.unbatched_query(octree, prefix, query_coords, blas_level, with_parents=True)\n",
    "print(pidx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"pyramid\" is a torch.int32 tensor of size (2, 4)\n",
      "Raw Content: \n",
      "[[1 3 3 0]\n",
      " [0 1 4 7]]\n",
      "\"pyramid\" summarizes the number of occupied \n",
      "cells per level, and their cumulative sum:\n",
      "\n",
      "Root node (implicitly defined):\n",
      "\thas 1 occupied cells\n",
      "\tstart idx (cumsum excluding current level): 0\n",
      "Level #1:\n",
      "\thas 3 occupied cells\n",
      "\tstart idx (cumsum excluding current level): 1\n",
      "Level #2:\n",
      "\thas 3 occupied cells\n",
      "\tstart idx (cumsum excluding current level): 4\n",
      "Final entry for total cumsum:\n",
      "\thas 0 occupied cells\n",
      "\tstart idx (cumsum excluding current level): 7\n"
     ]
    }
   ],
   "source": [
    "pyramid=pyramid_orig[0]\n",
    "describe_tensor(torch_tensor=pyramid, tensor_label='pyramid', with_shape=True, with_content=True)\n",
    "\n",
    "print('\"pyramid\" summarizes the number of occupied \\ncells per level, and their cumulative sum:\\n')\n",
    "for i in range(pyramid.shape[-1]):\n",
    "    if i ==0:\n",
    "        print(f'Root node (implicitly defined):')\n",
    "    elif i+1 < pyramid.shape[-1]:\n",
    "        print(f'Level #{i}:')\n",
    "    else:\n",
    "        print(f'Final entry for total cumsum:')\n",
    "    print(f'\\thas {pyramid[0,i]} occupied cells')\n",
    "    print(f'\\tstart idx (cumsum excluding current level): {pyramid[1,i]}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wisp",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
